# coding: utf-8
# date: 2022-11-02
# 基于 FR 模型，在 bdf 文件中设置阻尼值
import math
import os
import toml
import numpy
from matplotlib import pyplot


def readTOML():
    '''
    读取配置文件
    '''
    cfgTOML = 'cfg.toml'
    if not os.path.exists(cfgTOML):
        print('No TOML file found!')
        return
    with open(cfgTOML, 'r') as f:
        cfg = toml.load(f)

    param = cfg['parameter']
    w1 = param['omega1']
    w2 = param['omega2']
    w3 = param['omega3']
    wn = param['omegan']
    h = param['h']
    isplot = param['isplot']
    isbdf = param['isbdf']
    zone = param['zone']
    return w1, w2, w3, wn, h, isplot, isbdf, zone


def getK(w1, w2, w3, wn, h):
    '''
    计算斜率
    '''
    H = 2*math.pi*h
    k1 = H/(w1)
    k2 = H/(w2-w3)
    return k1, k2


if __name__ == '__main__':

    # 读取配置文件
    w1, w2, w3, wn, h, isplot, isbdf, zone = readTOML()
    print('CFG: ', w1, w2, w3, wn, h, isplot, isbdf, zone)

    if (not isbdf and not isplot):
        print('Nothing to do!')
        exit()

    # 计算斜率
    k1, k2 = getK(w1, w2, w3, wn, h)
    print('K: ', k1, k2)

    # 计算坐标
    if isplot:
        print('*** Plotting')
        x = numpy.arange(0.0, w3+1.0, 0.1)
        y = numpy.empty(len(x))
        for i in range(len(x)):
            if (x[i] <= w1):
                y[i] = k1*x[i]
            elif (w1 < x[i] and x[i] <= w2):
                y[i] = 2*math.pi*h
            elif (w2 < x[i] and x[i] <= w3):
                y[i] = 2*math.pi*h + k2*(x[i]-w2)
            else:
                y[i] = 0

        # 绘图
        pyplot.xlabel(r'$ \omega $ (rad)')
        pyplot.ylabel(r"$ {E_D}/{E_{S0}} $")
        pyplot.title('Frequency Range Model')
        pyplot.xlim(0, 20+1.0)
        pyplot.ylim(0, 2*math.pi*h+1.0)
        pyplot.plot(x, y)
        pyplot.show()

    elif isbdf and not isplot:
        print('*** Writing BDF')
        x = numpy.arange(zone[0], zone[1], zone[2])
        y = numpy.empty(len(x))
        for i in range(len(x)):
            if (x[i] <= w1):
                y[i] = k1*x[i]
            elif (w1 < x[i] and x[i] <= w2):
                y[i] = 2*math.pi*h
            elif (w2 < x[i] and x[i] <= w3):
                y[i] = 2*math.pi*h + k2*(x[i]-w2)
            else:
                y[i] = 0

        # 输出
        print('Omega: ', x)
        print('Range: ', y)

    print ('All done!')
